import util
from CR_DSPPytorch import decimate
import torch
import numpy as np
import matplotlib.pyplot as plt
import os

BASE_DIR = os.path.dirname(__file__)

fig_save_dir = os.path.join(BASE_DIR, 'show_figure', 'decimate')
if not os.path.exists(fig_save_dir):
    os.makedirs(fig_save_dir)

SNR_dB = 15

prbs = np.random.randint(low=0, high=2, size=(1, 131072))

sig_tx = util.map(prbs=prbs)
sig_tx = np.stack([sig_tx, sig_tx], axis=-1)
sig_tx = sig_tx.reshape(1, -1)

noise_var = 10**(-SNR_dB / 10)
noise = np.random.randn(*sig_tx.shape, 2) * np.sqrt(noise_var / 2)

sig = sig_tx + noise[..., 0] + 1j * noise[..., 1]

# plt.plot(sig.real,sig.imag, '.')
# plt.savefig(os.path.join(fig_save_dir, 'decimate.png'))
# plt.close()



sig_freq_wo_antialias = np.fft.fft(sig[..., 0::2])
plt.plot(np.fft.fftshift(np.abs(sig_freq_wo_antialias).squeeze()))
plt.title('1 samples per symbol without anti-alias filter')
plt.savefig(os.path.join(fig_save_dir, 'sig_freq_wo_antialias.png'))
plt.close()

sig_decimate = decimate(torch.from_numpy(
    np.stack([sig.real, sig.imag], axis=-1)),
                        q=2,
                        time=0,
                        est_best_time=True,
                        complex_out=True).data.cpu().numpy()
sig_decimate_f = np.fft.fft(sig_decimate)
plt.plot(np.fft.fftshift(np.abs(sig_decimate_f).squeeze()))
plt.title('1 samples per symbol with anti-alias filter')
plt.savefig(os.path.join(fig_save_dir, 'sig_freq_decimate.png'))
plt.close()


util.colorfulPlot(sig_decimate, prbs, saveFig=True, savePath=os.path.join(fig_save_dir, 'sig_const_decimate.png'))

SNR_measure_wo_decimate = util.effective_snr(sig[..., 0::2], prbs)
SNR_measure_decimate = util.effective_snr(sig_decimate, prbs)
print('SNR without decimate: {}'.format(SNR_measure_wo_decimate))
print('SNR with decimate: {}'.format(SNR_measure_decimate))
